
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import tushare as ts
import talib
from sklearn.preprocessing import StandardScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import classification_report, confusion_matrix, accuracy_score

plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号


# ==================== 1. 数据获取 ====================
def get_stock_data():
    """使用Tushare获取股票数据"""
    print("正在获取股票数据...")

    # 初始化pro接口
    pro = ts.pro_api('1546c65ff7671ac22882f792daabe8ee2943a82f5115fe252aa68296')

    # 获取日线数据
    df = pro.daily(
        ts_code="300750.SZ",
        start_date="20200101",
        end_date="20250414"
    )

    # 数据预处理
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values('trade_date').set_index('trade_date')
    df = df.rename(columns={
        'open': 'Open',
        'high': 'High',
        'low': 'Low',
        'close': 'Close',
        'vol': 'Volume'
    })

    print(f"获取到{len(df)}条数据")
    return df[['Open', 'High', 'Low', 'Close', 'Volume']]


# ==================== 2. 使用TA-Lib计算技术指标 ====================
def calculate_technical_indicators(df):
    """使用TA-Lib计算技术指标"""
    print("正在计算技术指标...")

    # 1. 趋势指标
    df['SMA_10'] = talib.SMA(df['Close'], timeperiod=10)
    df['EMA_12'] = talib.EMA(df['Close'], timeperiod=12)
    df['EMA_26'] = talib.EMA(df['Close'], timeperiod=26)

    # 2. MACD指标
    df['MACD'], df['MACD_Signal'], df['MACD_Hist'] = talib.MACD(
        df['Close'], fastperiod=12, slowperiod=26, signalperiod=9)

    # 3. 动量指标
    df['RSI'] = talib.RSI(df['Close'], timeperiod=14)
    df['STOCH_K'], df['STOCH_D'] = talib.STOCH(
        df['High'], df['Low'], df['Close'],
        fastk_period=14, slowk_period=3, slowd_period=3)

    # 4. 波动率指标
    df['ATR'] = talib.ATR(
        df['High'], df['Low'], df['Close'], timeperiod=14)
    df['BB_UPPER'], df['BB_MIDDLE'], df['BB_LOWER'] = talib.BBANDS(
        df['Close'], timeperiod=20, nbdevup=2, nbdevdn=2)

    # 5. 成交量指标
    df['OBV'] = talib.OBV(df['Close'], df['Volume'])
    df['ADOSC'] = talib.ADOSC(
        df['High'], df['Low'], df['Close'], df['Volume'], fastperiod=3, slowperiod=10)

    # 6. 其他指标
    df['ADX'] = talib.ADX(
        df['High'], df['Low'], df['Close'], timeperiod=14)
    df['CCI'] = talib.CCI(
        df['High'], df['Low'], df['Close'], timeperiod=14)

    # 删除含有NaN的行
    df = df.dropna()

    return df


# ==================== 3. 准备目标变量 ====================
def prepare_target(df, horizon=5):
    """准备目标变量"""
    print("正在准备目标变量...")

    # 计算未来horizon天的收益率
    df['future_return'] = df['Close'].pct_change(horizon).shift(-horizon)

    # 转换为分类问题：1=上涨，0=下跌
    df['target'] = (df['future_return'] > 0).astype(int)

    # 删除缺失值
    df = df.dropna()

    return df


# ==================== 4. 模型训练与评估 ====================
def train_and_evaluate(X, y):
    """训练并评估模型"""
    # 数据标准化
    scaler = StandardScaler()
    X_scaled = scaler.fit_transform(X)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        X_scaled, y, test_size=0.2, random_state=42, shuffle=False)

    # 训练随机森林
    print("正在训练随机森林模型...")
    rf = RandomForestClassifier(random_state=42)

    param_grid = {
        'n_estimators': [100, 200],
        'max_depth': [None, 5, 10],
        'min_samples_split': [2, 5]
    }

    grid_search = GridSearchCV(
        estimator=rf,
        param_grid=param_grid,
        cv=5,
        n_jobs=-1,
        verbose=1
    )
    grid_search.fit(X_train, y_train)

    best_rf = grid_search.best_estimator_
    print("最佳参数:", grid_search.best_params_)

    # 评估模型
    y_pred = best_rf.predict(X_test)

    print("\n混淆矩阵:")
    print(confusion_matrix(y_test, y_pred))

    print("\n分类报告:")
    print(classification_report(y_test, y_pred))

    print(f"\n准确率: {accuracy_score(y_test, y_pred):.4f}")

    # 特征重要性
    feature_importance = pd.DataFrame({
        'Feature': X.columns,
        'Importance': best_rf.feature_importances_
    }).sort_values('Importance', ascending=False)

    print("\n特征重要性:")
    print(feature_importance)

    # 可视化特征重要性
    plt.figure(figsize=(10, 6))
    plt.barh(feature_importance['Feature'][:10], feature_importance['Importance'][:10])
    plt.title('Top 10 重要特征')
    plt.xlabel('重要性')
    plt.tight_layout()
    plt.savefig('tezheng.png', dpi=300, bbox_inches='tight')
    plt.show()

    return best_rf, y_pred, X_test, y_test


# ==================== 5. 策略回测 ====================
def backtest_strategy(df, y_test, y_pred):
    """回测交易策略"""
    print("正在进行策略回测...")

    # 获取测试集日期
    test_dates = df.index[-len(y_test):]

    # 创建回测DataFrame
    backtest_df = pd.DataFrame({
        'Date': test_dates,
        'Actual': y_test,
        'Predicted': y_pred,
        'Close': df['Close'].values[-len(y_test):]
    }, index=test_dates)

    # 计算每日收益率
    backtest_df['Daily_Return'] = backtest_df['Close'].pct_change()

    # 策略收益率：当预测上涨时买入，否则不持仓
    backtest_df['Strategy_Return'] = backtest_df['Predicted'].shift(1) * backtest_df['Daily_Return']

    # 累计收益率
    backtest_df['Cumulative_Market'] = (1 + backtest_df['Daily_Return']).cumprod()
    backtest_df['Cumulative_Strategy'] = (1 + backtest_df['Strategy_Return']).cumprod()

    # 计算年化收益率和夏普比率
    annual_return = backtest_df['Strategy_Return'].mean() * 252
    annual_volatility = backtest_df['Strategy_Return'].std() * np.sqrt(252)
    sharpe_ratio = annual_return / annual_volatility

    print(f"\n年化收益率: {annual_return:.4f}")
    print(f"年化波动率: {annual_volatility:.4f}")
    print(f"夏普比率: {sharpe_ratio:.4f}")

    # 绘制收益率曲线
    plt.figure(figsize=(12, 6))
    plt.plot(backtest_df['Cumulative_Market'], label='买入持有策略')
    plt.plot(backtest_df['Cumulative_Strategy'], label='随机森林策略')
    plt.title('策略回测结果')
    plt.xlabel('日期')
    plt.ylabel('累计收益率')
    plt.legend()
    plt.grid()
    plt.savefig('shouyiquxian.png',dpi=300,bbox_inches='tight')
    plt.show()

    return backtest_df


# ==================== 主程序 ====================
if __name__ == "__main__":
    # 1. 获取数据
    df = get_stock_data()

    # 2. 计算技术指标
    df = calculate_technical_indicators(df)

    # 3. 准备目标变量
    df = prepare_target(df, horizon=5)

    # 4. 准备特征和目标
    X = df.drop(['future_return', 'target'], axis=1)
    y = df['target']

    # 5. 训练并评估模型
    model, y_pred, X_test, y_test = train_and_evaluate(X, y)

    # 6. 策略回测
    backtest_df = backtest_strategy(df, y_test, y_pred)